""" Implementation of helpers to parse DICOM format
Requirement:
    dicom(pydicom==0.9.9) is used, however the latest is 1.0
"""

from glob import glob
import logging

import numpy as np
import dicom as pydicom


def parse_dcm_dir(dicom_dir, order='InstanceNumber'):
    """
    Parse info from a directory of dicom files
    :param dicom_dir: target directory
    :param order: 'InstanceNumber', 'Location'
    :return:
        vol: 3D volume array
        sop: a list of sop uid
        loc: a list of locations
        spacing: 3-tuple pixel spacing
        slice_thickness: if Slice_Thickness in dicom is available. It might be inconsistent with locations
        series_uid: series uid for identification
    """

    dicom_names = glob('%s/*.dcm' % dicom_dir)  # get all the dicom files in the dicom_dir
    slices_list = list()

    for dicom_name in dicom_names:
        slices_list.append(parse_single_dcm(dicom_name)) # parse info from a single dicom file

    if order == 'InstanceNumber':
        slices_list.sort(key=(lambda x: int(x[1].InstanceNumber)))
    elif order == 'Location':
        slices_list.sort(key=(lambda x: get_location(x[1])))
        mean_location = np.mean([get_location(dicom) for _, dicom in slices_list])
        reverse_flag = bool(mean_location <= 0)
        if reverse_flag:
            slices_list = slices_list[::-1]
            logging.debug('non positive mean location; then reverse')
    else:
        raise ValueError('unknown order')

    vol = list()
    sop = list()
    loc = list()
    series_uid = None
    spacing = None
    slice_thickness = None

    for slice, dicom in slices_list:
        vol.append(slice)
        sop.append(str(dicom.SOPInstanceUID))
        loc.append(get_location(dicom))

        # check series_uid
        if series_uid == str(dicom.SeriesInstanceUID):
            pass
        elif series_uid is None:
            series_uid = str(dicom.SeriesInstanceUID)
        else:
            logging.error('series_uid inconsistent: %s and %s' %(series_uid, str(dicom.SeriesInstanceUID)))

        # check spacing
        if spacing is None:
            spacing = np.asarray(dicom.PixelSpacing)
        elif np.max(spacing - np.asarray(dicom.PixelSpacing)) > 0.1:
            logging.error('spacing inconsistent')

        # try to get slice_thickness
        if slice_thickness is None:
            slice_thickness = dicom.get('SliceThickness')
        elif slice_thickness != dicom.get('SliceThickness'):
            slice_thickness = -1

    vol = np.transpose(np.asarray(vol), (1, 2, 0)).astype(np.int16)

    # adjust slice thickness and spacing
    if slice_thickness is None:
        logging.warning('SliceThickness missing')
        slice_thickness = None
    if slice_thickness == -1:
        logging.warning('SliceThickness inconsistent')
        slice_thickness = None

    emp_diff = np.fabs(loc[2] - loc[3])
    sum_diff = []
    for ind in range(len(loc) - 1):
        diff = np.fabs(loc[ind] - loc[ind + 1])
        sum_diff.append(diff)

    mean_diff = np.mean(sum_diff)
    if slice_thickness is not None and np.fabs(slice_thickness - mean_diff) < 0.1:
        spacing = np.append(spacing, slice_thickness)
    else:
        spacing = np.append(spacing, emp_diff)

    return vol, sop, loc, spacing, slice_thickness, series_uid


def parse_single_dcm(dicom_name):
    """
    Parse info from a single dicom
    :param dicom_name: target dicom file
    :return:
        slice: 2D numpy array
        dicom: a class-like struct
    """
    dicom = pydicom.read_file(dicom_name)  # pydicom return a class-like struct

    # rescale intercept and slope
    slice = dicom.pixel_array
    try:
        intercept = dicom.RescaleIntercept
        slope = dicom.RescaleSlope
    except AttributeError:
        intercept = 0
        slope = 1
        logging.warning('rescale info missing.')
    finally:
        slice = slice * slope + intercept

    return slice, dicom


# helpers
def get_location(dicom):
    """
    Get location info from a dicom. Try to get SliceLocation first. If not exist, then ImagePositionPatient.
    :param dicom: dicom info returned by pydicom.read_file
    :return:
        location: scalar
    """
    location = None

    try:
        location = np.float(dicom.SliceLocation)
    except AttributeError:
        logging.warning('dicom miss location info.')
        try:
            location = np.float(dicom.ImagePositionPatient[2])
        except AttributeError:
            logging.warning('dicom miss image position.')

    return location


def check_dcm_dir(dicom_dir):
    dicom_names = glob('%s/*.dcm' % dicom_dir)

    if not dicom_names:
        logging.debug('%s contains no dicoms.' % dicom_dir)
        return False

    dicom = pydicom.read_file(dicom_names[0], stop_before_pixels=True)
    
    if str(dicom.Modality) != 'CT':
        logging.debug('%s is not a CT directory.' % dicom_dir)
        return False

    return True


def check_location(location, slice_thickness):
    max_diff = 0.5 * slice_thickness
    min_diff = 1e-2

    for ind in range(len(location)-1):
        diff = np.fabs(location[ind] - location[ind+1])
        if diff < min_diff:
            return 'too small diff: (%f, %f). slice_thickness = %f' \
                   % (location[ind], location[ind + 1], slice_thickness)
        elif np.fabs(diff-slice_thickness) >= max_diff:
            return 'too large diff: (%f, %f). slice_thickness = %f' \
                   % (location[ind], location[ind + 1], slice_thickness)

    return None


